import psycopg2

from django.db.utils import load_backend,DatabaseError
from .introspection import TenantSchemaDatabaseIntrospection
# https://gitlab.com/-/snippets/2137539
psql_backend = load_backend('django.db.backends.postgresql')


class DatabaseWrapper(psql_backend.DatabaseWrapper):
    def __init__(self,*args,**kwargs) -> None:
        super().__init__(*args,**kwargs)
    

        self.schemas = ('public',)
        self.introspection = TenantSchemaDatabaseIntrospection(self)
        self.is_search_path_set = False 


    @property
    def primary_schema(self):
        return self.schemas[0]


    def close(self):
        # 无聊函数
        super().close()
        self.is_search_path_set = False


    def rollback(self):
        super().rollback()
        self.is_search_path_set = False


    def set_schema(self,*schemas,include_public=True):
        # positional argument
        self.schemas = schemas

        if include_public and 'public' not in schemas:
            self.schemas +=('public',)

        self.is_search_path_set = False
        # keyword argument

    def _cursor(self,name=None):
        """
        所有django发出的数据库操作指令，都需要先调用此函数以获得数据库的操作对象
        因而我们提前地加上sert search_path 命令，保证所有的操作都指向当前锁定的schema
        """
        cursor = super()._cursor(name=name)

        if self.is_search_path_set:
            return cursor
        try:
            target_path = ','.join(self.schemas)
            search_path_str= f'SET search_path = {target_path}'
            cursor.execute(search_path_str)
        except (DatabaseError,psycopg2.InterfaceError):
            self.is_search_path_set = False
        else:
            self.is_search_path_set = True

        return cursor